#!/usr/bin/env python2.5

"""
A test OpenID consumer and provider based on JanRain's Python OpenID library,
version 2.1.1, from http://www.openidenabled.com/python-openid/.

See http://code.google.com/p/openid-test/ and
http://code.google.com/p/openid-test/wiki/TestSpec.

Responds to the following URLs:

  .../caps
  .../<protocol>/identity/will-sign
  .../<protocol>/identity/will-setup-then-sign
  .../<protocol>/identity/will-setup-then-cancel
  .../<protocol>/rp

where <protocol> is "1.1" or "2.0" or "1.1,2.0". Make sure your realm has
"-test" or "openid-test" in its path.

To install, make openid_1_0_0 and openid_2_1_1 symlinks, pointing to the
respective versions of the janrain openid librari, in the same directory as
this file. Then, configure your web server to run this file as a CGI script,
which usually requires the executable bit to be set.

For more about OpenID, see:
  http://openid.net/
  http://openid.net/about.bml

Copyright 2008 Ryan Barrett
"""

__author__ = 'Ryan Barrett <openid-test@ryanb.org>'

import cgitb; cgitb.enable()

import BaseHTTPServer
import cgi
import Cookie
import cPickle
import datetime
import logging
import os
import pprint
import re
import sys
import traceback
import urllib
import urlparse
import wsgiref.handlers
import wsgiref.util

# these are the openid.server.server, openid.consumer.consumer, and
# openid.store.filestore modules. they're imported from openid_1_0_0 or
# openid_2_1_1 in main(), depending on the version in the URL.
server = None
consumer = None
filestore = None
dumbstore = None

logging.getLogger().addHandler(logging.StreamHandler())
logging.getLogger().setLevel(logging.DEBUG)

# supported openid protocol versions
PROTOCOL_VERSIONS = ('1.1', '2.0', '1.1,2.0')

# name of session argument for 1.0.0 library's token
TOKEN_ARG = '_openid_consumer_.last_token'

# TODO: remove
# name of query parameter for the http method; used when emulating the 2.1.1
# library with 1.0.0.
# HTTP_METHOD_ARG = '__http_method'

# trust root (realm) to send to provider when acting as consumer
TRUST_ROOT = 'http://' + os.environ['SERVER_NAME']
if os.environ['SERVER_PORT'] != '80':
  TRUST_ROOT += ':' + os.environ['SERVER_PORT']

# base url, usually the URI path to this script with a trailing /
if 'SCRIPT_NAME' in os.environ:
  BASE_URL = TRUST_ROOT + os.environ['SCRIPT_NAME'] + '/'
else:
  # default
  BASE_URL = TRUST_ROOT + '/sandbox/cgi-bin/janrain_python.cgi/'


# regex used to parse URL paths
PATH_REGEX = re.compile(
  """(caps |
     (%s)/
       (identity/will-sign |
        identity/will-setup-then-sign |
        identity/will-setup-then-cancel |
        rp |
        rp\.return_to)
     )$
     """ % '|'.join(PROTOCOL_VERSIONS),
  re.VERBOSE)


# HTML templates
DOCTYPE = """<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Strict//EN"
"http://www.w3.org/TR/xhtml1/DTD/xhtml1-strict.dtd">
"""

HEADER = """<html><head>
<link rel="openid.server" href="%%(uri)s" />
<base href="%s" /> 
</head>
<body>
""" % BASE_URL

FOOTER = '</body></html>'

ERROR = """
<p>Error: %(message)s</p>
<p>Request: %(request)s</p>
<p>Form: %(form)s</p>
"""

ENDPOINT = """
<p>Hi! This is a test OpenID endpoint. You can use this URL as an OpenID
identifier in a OpenID consumer site:</p>
<p><i><a href="%(uri)s">%(uri)s</a></i></p>
"""

CONSUMER_FORM = """
<form method='get'>
  <label for='openid_identifier'> OpenID URL to test: </label>
  <input type='text' size=50 name='openid_identifier' />
  <input type='submit' value='Go' />

  <table><tr><td>
    Operation:<br />
    <input type='radio' name='op' value='' id='op_everything' checked />
      <label for='op_everything'>Everything (full OpenID flow)</label><br />
    <input type='radio' name='op' value='disco' id='op_disco' />
      <label for='op_disco'>Return discovery information</label><br />
    <input type='radio' name='op' value='assoc' id='op_assoc' />
      <label for='op_assoc'>Return association details</label>
  </td><td>
    Association mode:<br />
    <input type='radio' name='assoc_mode' value='stateful' id='stateful' checked />
      <label for='stateful'>Stateful, possibly cached</label><br />
    <input type='radio' name='assoc_mode' value='stateful_new' id='stateful_new' />
      <label for='stateful_new'>Stateful, create new association</label><br />
    <input type='radio' name='assoc_mode' value='stateless' id='stateless' />
      <label for='stateless'>Don't associate (dumb mode)</label><br />
    <input type='radio' name='assoc_mode' value='use_handle' id='use_handle' />
      <label for='use_handle'>Use association handle:</label>
    <input type='text' name='assoc_handle' />
  </td></tr></table>
</form>
"""

OK = 'OK\n'
CANCELLED = 'Cancelled\n'
SETUP_NEEDED = 'Setup needed\n'

DISCOVERY = """
user_specified_url: %(claimed_id)s
openid_provider: %(openid_provider)s
local_id: %(local_id)s
"""

ASSOCIATION = """
handle: %(handle)s
type: %(type)s
"""

FRONT_PAGE = """
Welcome to the
<a href="http://www.openidenabled.com/python-openid/">JanRain Python library</a>'s
OpenID interopability test server. I'm part of the
<a href="http://code.google.com/p/openid-test/">openid-test project</a>.
<ul><li>Capabilities:<ul>
<li><a href="caps">caps</a></li>
</ul></li>
<li>Relying party (consumer) URLs:<ul>
%s
</ul></li>
<li>Identity URLs which return signatures right away:<ul>
%s
</ul></li>
<li>Identity URLs which fail, but after setup succeed:<ul>
%s
</ul></li>
<li>Identity URLs which fail, and then after setup do openid.mode=cancel:<ul>
%s
</ul></li>
</ul>
""" % tuple(
  ['\n'.join([template % (ver, ver) for ver in PROTOCOL_VERSIONS])
   for template in [
     '<li><a href="%s/rp">%s/rp</a></li>',
     '<li><a href="%s/identity/will-sign">%s/identity/will-sign</a></li>',
     '<li><a href="%s/identity/will-setup-then-sign">%s/identity/will-setup-then-sign</a></li>',
     '<li><a href="%s/identity/will-setup-then-cancel">%s/identity/will-setup-then-cancel</a></li>',
     ]])

CAPS = """
# Based on JanRain's Python OpenID library, version 2.1.1.
# See http://www.openidenabled.com/python-openid/

openid1.1
openid2.0
# xri
"""

class Handler:
  """Base request handler class."""
  def __init__(self, environ, start_response):
    self.environ = environ
    self.uri = wsgiref.util.request_uri(self.environ, include_query=False)
    self.request = wsgiref.util.request_uri(self.environ, include_query=True)
    self.start_response = start_response
    self.output = []

  def arg_dict(self):
    """Flattens the URL parameters and form data into a dictionary.

    Only includes the first value for each parameter.
    """
    form = cgi.FieldStorage()
    return dict([(arg, form.getfirst(arg)) for arg in form])

  def get_openid_request(self):
    """Creates an OpenIDRequest for this request, if appropriate.

    If this request is not an OpenID request, returns False. If an error occurs
    while parsing the arguments, renders an error page and returns None. (Ugh.
    I know.)
    """

    try:
      oidrequest = self.openid_server.decodeRequest(self.arg_dict())
      if oidrequest:
        logging.debug('Received OpenID request: %s' % oidrequest)
      else:
        return False
      return oidrequest
    except:
      trace = ''.join(traceback.format_exception(*sys.exc_info()))
      self.render(ERROR, 'Error parsing OpenID request:\n%s' % trace, '400 Error')
      return None

  def respond(self, oidresponse):
    """Send an OpenID response.
    """
    logging.debug('Sending OpenID response: %s' % oidresponse)
    encoded_response = self.openid_server.encodeResponse(oidresponse)

    # build the HTTP status code string
    message = BaseHTTPServer.BaseHTTPRequestHandler.responses.get(
      encoded_response.code, '')
    status_str = '%s %s' % (encoded_response.code, message)

    headers = encoded_response.headers.items()
    headers.append(('Content-Type', 'text/plain'))
    print >> sys.stderr, headers
    self.start_response(status_str, headers)

    logging.debug('Sending encoded response: %s' % encoded_response)
    if encoded_response.body:
      logging.debug('Sending response body: %s' % encoded_response.body)
      self.output = [encoded_response.body]

  def render(self, template, message='', status_code='200 OK',
             content_type='text/html', headers=[], data={}):
    """Render the given template, including the extra (optional) values.
    """
    self.start_response(status_code, [('Content-Type', content_type)] + headers)
    form = cgi.FieldStorage()
    data_dict = {
      'uri': self.uri,
      'request': self.request,
      'form': pprint.pformat(form),
      'message': message,
    }
    data_dict.update(self.arg_dict())
    data_dict.update(data)

    if content_type == 'text/html':
      self.output = [DOCTYPE, HEADER % data_dict, template % data_dict, FOOTER]
    else:
      self.output = [template % data_dict]

    # work around a bug in wsgiref.handlers that makes it not support unicode.
    self.output = [str(part) for part in self.output]

  def get(self):
    self.render(ENDPOINT, 'Welcome!')


class Error(Handler):
  """Displays an error."""
  def __init__(self, environ, start_response):
    Handler.__init__(self, environ, start_response)

  def get(self):
    self.render(ERROR, '404 Not found', '404 Not found')

  post = get


class FrontPage(Handler):
  """Displays the front page with all of the supported links."""
  def get(self):
    self.render(FRONT_PAGE)


class Capabilities(Handler):
  """Displays the capabilities page."""
  def get(self):
    self.render(CAPS, content_type='text/plain')


class Provider(Handler):
  """The base OpenID provider request handler class."""
  def __init__(self, environ, start_response, checkid_immediate=True,
               checkid_setup=True):
    """The keyword argument values determine how we respond to each mode."""
    Handler.__init__(self, environ, start_response)
    self.checkid_immediate = checkid_immediate
    self.checkid_setup = checkid_setup
    self.openid_server = server.Server(filestore.FileOpenIDStore('.'), self.uri)

  def post(self):
    """Handles associate and check_authentication requests."""
    oidrequest = self.get_openid_request()
    if oidrequest and oidrequest.mode in ('associate', 'check_authentication'):
      self.respond(self.openid_server.handleRequest(oidrequest))
    elif oidrequest is False:
      msg = 'Expected associate or check_authentication, got %s' % oidrequest.mode
      self.render(ERROR, msg, '404 Not found')
    # else get_openid_request() rendered the error page

  def get(self):
    """Handles checkid_* requests."""
    oidrequest = self.get_openid_request()
    if oidrequest and oidrequest.mode in ('checkid_immediate', 'checkid_setup'):
      self.respond(oidrequest.answer(getattr(self, oidrequest.mode)))
    elif oidrequest is False:
      Handler.get(self)
    # else get_openid_request() rendered the error page


class Sign(Provider):
  """Signs a request and verifies the signature."""
  def __init__(self, environ, start_response):
    """The keyword argument values determine how we respond to each mode."""
    Provider.__init__(self, environ, start_response, checkid_immediate=True,
                      checkid_setup=True)


class SetupThenSign(Sign):
  """Does setup for a new request, then signs and verifies a signature."""
  def __init__(self, environ, start_response):
    """The keyword argument values determine how we respond to each mode."""
    Provider.__init__(self, environ, start_response, checkid_immediate=False,
                      checkid_setup=True)


class SetupThenCancel(Sign):
  """Does setup for a new request, then signs it and sends a cancel back."""
  def __init__(self, environ, start_response):
    """The keyword argument values determine how we respond to each mode."""
    Provider.__init__(self, environ, start_response, checkid_immediate=False,
                      checkid_setup=False)


class Consumer(Handler):
  """The base consumer handler."""
  def __init__(self, *args):
    Handler.__init__(self, *args)
    arg_dict = self.arg_dict()

    # set up the association store based on the assoc_mode query parameter
    store = filestore.FileOpenIDStore('.')

    assoc_mode = arg_dict.get('assoc_mode')
    if assoc_mode:
      if assoc_mode == 'stateful':
        pass
      elif assoc_mode in ('stateful_new', 'use_handle'):
        # remove the other associations
        for file, assoc in store._allAssocs():
          if not (assoc_mode == 'use_handle' and
                  assoc.handle == arg_dict.get('assoc_handle')):
            filestore._removeIfPresent(file)
        store = filestore.FileOpenIDStore('.')
      elif assoc_mode == 'stateless':
        store = dumbstore.DumbStore('passphrase')
      else:
        assert False, 'bad assoc_mode value: %s' % assoc_mode

    # retrieve the session from the cookie, if available. obviously not secure.
    session = self.get_session_from_cookie()
    self.openid_consumer = consumer.Consumer(session, store)

  def make_session_cookie(self):
    """Pickles the session and returns it as a cookie value."""
    expires = datetime.datetime.now() + datetime.timedelta(minutes=5)
    expires_rfc822 = expires.strftime('%a, %d %b %Y %H:%M:%S -0800')
    encoded = urllib.quote(cPickle.dumps(self.openid_consumer.session))
    return 'janrain-python-session=%s; expires=%s' % (encoded, expires_rfc822)

  def get_session_from_cookie(self):
    """Recovers the user's session from their cookie and returns it. If they
    don't have a session cookie, returns None."""
    cookies = os.environ.get('HTTP_COOKIE')
    if cookies:
      morsel = Cookie.BaseCookie(cookies).get('janrain-python-session')
      if morsel:
        session = cPickle.loads(urllib.unquote(morsel.value))
        assert isinstance(session, dict)
        return session

    return {}

  def get(self):
    """Handles consumer GET requests."""
    args = self.arg_dict()

    if 'openid.mode' in args:
      # this is a redirect from a provider. attempt to verify the login.
      response = self.openid_consumer.complete(args, self.uri)
      if response.status == 'success':
        self.render(OK, content_type='text/plain')
      elif response.status == 'cancel':
        self.render(CANCELLED, content_type='text/plain')
      elif response.status == 'setup_needed':
        self.render(SETUP_NEEDED, content_type='text/plain')
      else:
        self.render(ERROR, 'Provider responded: %s %s' %
                    (response.status, response.message))

    elif 'openid_identifier' in args:
      # start with discovery
      auth_request = self.openid_consumer.begin(args['openid_identifier'])

      # *optionally* stores the session in a cookie. if the client doesn't
      # support cookies, no matter, the consumer flow will still work.
      headers = [('Set-Cookie', self.make_session_cookie())]

      # this is a form post. what to do...?
      if args.get('op') == 'disco':
        # return the discovery details
        claimed_id = auth_request.endpoint.claimed_id
        local_id = auth_request.endpoint.local_id
        if local_id == claimed_id:
          local_id = ''

        discovered = {'claimed_id': claimed_id,
                      'openid_provider': auth_request.endpoint.server_url,
                      'local_id': local_id}
        self.render(DISCOVERY, content_type= 'text/plain', data=discovered,
                    headers=headers)
        return

      elif args.get('op') == 'assoc':
        # return the association we would have used
        if auth_request.assoc:
          assoc = {'handle': auth_request.assoc.handle,
                   'type': auth_request.assoc.assoc_type }
        else:
          assoc = {'handle': '', 'type': ''}

        self.render(ASSOCIATION, content_type='text/plain', data=assoc,
                    headers=headers)

      else:
        # start the flow!
        return_to = self.uri + '.return_to'
        redirect_to = auth_request.redirectURL(TRUST_ROOT, return_to)
        headers += [('Location', redirect_to)]
        self.start_response('302 Found', headers)

    else:
      # this is an initial GET. show the form.
      self.render(CONSUMER_FORM, 'unused')
      return


def app(environ, start_response):
  """Pass the request to the right handler."""
  handlers = {'caps': Capabilities,
              'identity/will-sign': Sign,
              'identity/will-setup-then-sign': SetupThenSign,
              'identity/will-setup-then-cancel': SetupThenCancel,
              'rp': Consumer,
              'rp.return_to': Consumer,
              }

  uri = wsgiref.util.request_uri(environ, include_query=False)
  match = PATH_REGEX.search(uri)
  method = environ['REQUEST_METHOD']

  # default to not found
  handler_cls = FrontPage

  if match and method in ('GET', 'POST'):
    # TODO(ryanb); ugh, this is awful. make this not awful.
    caps, openid_version, action = match.groups()

    if caps == 'caps':
      action = 'caps'
    else:
      import_libs(openid_version)
    if action in handlers:
      handler_cls = handlers[action]

  handler = handler_cls(environ, start_response)
  getattr(handler, method.lower())()
  return handler.output


def import_libs(openid_version):
  """Import the the JanRain libraries for the given OpenID version.

  The API changed between 1.0.0 and 2.1.1, so if this imports 1.0.0, it also
  fixes it up with a compatibility wrapper. This lets the rest of the code
  treat it like 2.1.1, regardless of which version it actually is.

  The openid_version arg should be '1.1', '2.0', or '1.1,2.0'.
  """
  global server
  global consumer
  global filestore
  global dumbstore

  # generally useful, easily initialized struct class
  class Struct:
    def __init__(self, **kwargs):
      for arg, val in kwargs.items():
        setattr(self, arg, val)

  # point to the requested version of the janrain libraries
  if openid_version in ('2.0', '1.1,2.0'):
    import openid_2_1_1
    sys.modules['openid'] = openid_2_1_1
    from openid_2_1_1.server import server
    from openid_2_1_1.consumer import consumer
    from openid_2_1_1.store import filestore

    # dumbstore was removed in 2.1.1. add a shim.
    dumbstore = Struct(DumbStore=lambda passphrase: None)

  else:
    assert openid_version == '1.1'
    import openid_1_0_0
    sys.modules['openid'] = openid_1_0_0
    from openid_1_0_0.server import server
    from openid_1_0_0.consumer import consumer
    from openid_1_0_0.store import filestore
    from openid_1_0_0.store import dumbstore

    # OpenIDConsumer => Consumer, OpenIDServer => Server
    consumer.Consumer = consumer.OpenIDConsumer
    server.Server = server.OpenIDServer

    # Consumer.beginAuth() => Consumer.begin()
    def consumer_begin(self, *args, **kwargs):
      status, auth_req = consumer.Consumer.beginAuth(self, *args, **kwargs)
      assert auth_req, status
      self.session[TOKEN_ARG] = auth_req.token
      auth_req.endpoint = auth_req
      auth_req.local_id = ''
      auth_req.claimed_id = auth_req.server_id
      auth_req.assoc = self.store.getAssociation(auth_req.server_url)
      auth_req.redirectURL = (lambda trust_root, return_to:
        self.constructRedirect(auth_req, return_to, trust_root))
      return auth_req
    consumer.Consumer.begin = consumer_begin

    # Consumer.completeAuth() => Consumer.complete()
    def complete(self, args, uri):
      status, message = self.completeAuth(self.session[TOKEN_ARG], args)
      if status == consumer.SUCCESS:
        if message:
          status = 'success'
        else:
          status = 'cancel'
      elif status == consumer.SETUP_NEEDED:
        status = 'setup_needed'
      return Struct(status=status, message=message)
    consumer.Consumer.complete = complete

    # Consumer.__init__() added session arg
    orig_consumer_init = consumer.Consumer.__init__
    def consumer_init(self, session, *args, **kwargs):
      orig_consumer_init(self, *args, **kwargs)
      self.session = session
    consumer.Consumer.__init__ = consumer_init

    # Consumer.__init__()'s positional args swapped. (wtf! i mean, really.)
    orig_server_init = server.Server.__init__
    server.Server.__init__ = (lambda self, store, endpoint:
                                orig_server_init(self, endpoint, store))

    # no separate decodeRequest() step or Request class in 1.0.0; fabricate one
    class Request:
      def __init__(self, server, args):
        self.server = server
        self.args = args
        self.mode = args['openid.mode']

      def answer(self, allow):
        return self.server.handleRequest(self, allow)

    def decodeRequest(self, args):
      if args.get('openid.mode'):
        return Request(self, args)
      else:
        return None
    server.Server.decodeRequest = decodeRequest

    # Server.getOpenIDResponse() => server.handleRequest()
    def handleRequest(self, req, allow=True):
      method = os.environ['REQUEST_METHOD']
      status, val = self.getOpenIDResponse(method, req.args, lambda x, y: allow)
      assert status != server.LOCAL_ERROR, val

      responses = {server.REDIRECT: (302, val),
                   # see 1.0.0 server.py:461
                   server.DO_AUTH: (200, val),
                   server.DO_ABOUT: (200, 'DO_ABOUT'),
                   server.REMOTE_OK: (200, val),
                   server.REMOTE_ERROR: (400, val),
                   }
      code, body = responses[status]
      if status == server.DO_AUTH:
        if allow:
          body = req.args['openid.return_to']
        else:
          body = val.getCancelURL()

      headers = {}
      if code == 302:
        headers['Location'] = body
      return Struct(code=code, body=body, headers=headers)

    server.Server.handleRequest = handleRequest

    # ignore encoding for now
    server.Server.encodeResponse = lambda self, response: response


def main(argv):
  """FUCKING GET ME A TOWEL!!!"""
  wsgiref.handlers.CGIHandler().run(app)

if __name__ == '__main__':
  main(sys.argv)
